#!/usr/bin/env python
# coding:utf-8
import numpy as np
import math
import argparse
import logging
import time
import cupy as cp
from typing import List, Tuple
import multiprocessing
from concurrent.futures import ProcessPoolExecutor
import os
from flask import Flask, request, jsonify
import t_datas
import threading

app = Flask(__name__)
dst_features_gpu = None
mysql_result_cpu = None

class GPUWorker:
    def __init__(self, gpu_id: int):
        self.gpu_id = gpu_id
        cp.cuda.Device(gpu_id).use()

    def process_batch(self, src_feature: np.ndarray, batch_features: List[np.ndarray],
                     start_idx: int, threshold: float) -> List[Tuple[int, float]]:
        """在指定GPU上处理一批特征"""
        #cp.cuda.Device(self.gpu_id).use()
        try:
            threshold = float(threshold)
            # 将源特征转移到GPU
            src_feature_gpu = cp.array(src_feature)
            src_norm = float(cp.linalg.norm(src_feature_gpu))
            # 将目标特征转移到GPU
            #dst_features_gpu = cp.array(batch_features)
            # GPU计算
            dot_products = cp.dot(batch_features, src_feature_gpu)
            dst_norms = cp.linalg.norm(batch_features, axis=1)
            cosine_similarities = dot_products / (dst_norms * src_norm)
            # 将结果转回CPU
            cosine_similarities_cpu = cp.asnumpy(cosine_similarities)

            # 清理GPU内存
            mempool = cp.get_default_memory_pool()
            pinned_mempool = cp.get_default_pinned_memory_pool()
            mempool.free_all_blocks()
            pinned_mempool.free_all_blocks()
            # 返回高于阈值的结果
            results = []
            for j, score in enumerate(cosine_similarities_cpu):
                if score >= threshold:
                    results.append((start_idx + j, float(score)))

            return results

        except Exception as e:
            logging.error(f"GPU {self.gpu_id} error: {str(e)}")
            return []

@app.route('/proc1', methods=['POST'])
def proc():
    global dst_features_gpu
    global mysql_result_cpu

    try:
        global_start_time = time.time()

        # 从请求中获取JSON数据
        data = request.json

        # 提取参数
        feature_src = data.get('feature')
        scheme = data.get('scheme')
        threshold = data.get('score')
        log_id = data.get('log_id')
        logging.info(f"{log_id}_______开始计时")

        if not mysql_result_cpu:
            #logging.error("No destination features found in file")
            mysql_result_cpu = t_datas.getTable(scheme, log_id)
        #mysql_result_cpu=datas.getTable(scheme,log_id)
        #mysql_results = [dict(zip(columns, row)) for row in mysql_results]
        mysql_time=time.time()
        logging.info(f"{log_id}___获取mysql时间：{mysql_time-global_start_time},共：{len(mysql_result_cpu)}")


        # 获取可用的GPU数量
        num_gpus = cp.cuda.runtime.getDeviceCount()

        # 解码源特征
        src_feature_arr = t_datas.decode_single_feature(feature_src)

        # 将目标特征批量传输到 GPU
        if dst_features_gpu is None:
            update_dst_features_gpu()
        dst_features_gpu_time=time.time()


        # 计算每个GPU处理的批次大小
        total_features = len(dst_features_gpu)

        logging.info(f"{log_id}___获取decode时间：{dst_features_gpu_time - mysql_time},共：{total_features}")
        batch_size = math.ceil(total_features / num_gpus)
        # 准备每个GPU的任务
        tasks = []
        for gpu_id in range(num_gpus):
            start_idx = gpu_id * batch_size
            end_idx = min(start_idx + batch_size, total_features)
            logging.info(f"{log_id}___{gpu_id}开始干活:{start_idx} , {end_idx}")
            if start_idx < end_idx:
                tasks.append((gpu_id, src_feature_arr, dst_features_gpu[start_idx:end_idx], start_idx, threshold))

        # 创建进程池进行并行处理
        results = []
        with ProcessPoolExecutor(max_workers=num_gpus) as executor:
            futures = []
            for gpu_id, src_feat, batch_feat, start_idx, thresh in tasks:
                worker = GPUWorker(gpu_id)
                future = executor.submit(worker.process_batch, src_feat, batch_feat, start_idx, thresh)
                futures.append(future)

            # 收集结果
            for future in futures:
                results.extend(future.result())

        end_time = time.time()
        total_time = end_time - global_start_time

        all_result=[]
        if not results:
            #logging.error("{}")
            return jsonify([])

        for result in results:
            id_value = mysql_result_cpu[result[0]][0]
            single = {"id": id_value, "score": result[1]}
            all_result.append(single)

        results=sorted(all_result, key=lambda x: x['score'],reverse=True)


        logging.info(f"{log_id}___Found {len(results)} matches:")

        global_end_time = time.time()
        total_time = global_end_time - global_start_time
        logging.info(f"{log_id}___总计算时间: {total_time:.4f}s")
        logging.info(f"{log_id}___Average processing speed: {len(mysql_result_cpu)/total_time:.2f} features/s")
        return jsonify(results)
    except Exception as e:
        logging.error(f"{log_id}___Error occurred: {e}")
        return jsonify([])


def update_dst_features_gpu():
    try:
        global dst_features_gpu,mysql_result_cpu
        log_id=time.time()
        logging.info(f"{log_id}___更新GPU")
        mysql_result_cpu=t_datas.getTable("album",log_id)
        dst_features = []
        for i, feature_dst in enumerate(mysql_result_cpu):
            try:
                dst_feature_arr = t_datas.decode_single_feature(feature_dst[1])
                dst_features.append(dst_feature_arr)
            except Exception as e:
                logging.warning(f"Failed to decode feature at index {i}: {e}")
        dst_features_gpu = cp.array(dst_features)
        logging.info(f"{log_id}___更新GPU完成 共{len(dst_features)}个")
    except Exception as e:
        logging.error(f"Failed to update dst_features_gpu: {e}")

if __name__ == "__main__":
    logging.basicConfig(
        filename= os.path.expanduser('~')+'/app_.log',
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )
    multiprocessing.set_start_method("spawn")
    updater_thread = threading.Thread(target=update_dst_features_gpu,  daemon=True)
    updater_thread.start()
    app.run(host='0.0.0.0', port=6001)